Add quantization support for DiffusionGemma#1935
Conversation
Signed-off-by: lvliang-intel <liang1.lv@intel.com>
…upport_diffusiongemma
for more information, see https://pre-commit.ci
|
/azp run Unit-Test-CUDA-AutoRound |
|
Azure Pipelines successfully started running 1 pipeline(s). |
Transformers backend test result:test_diffusiongemma_transformers.py CUDA_VISIBLE_DEVICES=2 python test_diffusiongemma_transformers.py [1/5] Loading config ... [2/5] Loading tokenizer ... [3/5] Loading generation_config ... [4/5] Loading quantized model onto cuda:0 ... [5/5] Running a simple Q&A ... Q: Why is the sky blue? Please answer in one short paragraph. transformers load + inference succeeded |
…/auto-round into lvl/support_diffusiongemma
|
/azp run Unit-Test-CUDA-AutoRound |
|
Azure Pipelines successfully started running 1 pipeline(s). |
|
|
||
| if modules_to_not_convert and not user_specified: | ||
| _DEFAULT_SKIP_KEYWORDS = ("embed", "embed_tokens", "lm_head", "output_embed", "norm") | ||
| modules_to_not_convert = [ |
There was a problem hiding this comment.
Any large change in this file is risky.
Could you provide more details? From my understanding, everything listed above except lm_head should no longer be a linear layer and therefore should already be handled correctly. As for lm_head, it has its own dedicated logic, so I'm not sure why changes would be needed there.
There was a problem hiding this comment.
The original logic is correct for normal models, the get_keys_to_not_convert returns only "lm_head/embed", which aren't quantizable linear layers, so a filter would be a no-op. But DiffusionGemma breaks that assumption because of one regex line in transformers "modeling_diffusion_gemma.py" _tied_weights_keys.
r"encoder.language_model.layers.(?:[^.]+.)*weight" → r"decoder.layers.(?:[^.]+.)*weight"
This ties every .weight under the encoder layers to the matching decoder weight, i.e. self_attn.q_proj.weight, k_proj.weight, v_proj.weight, o_proj.weight, and the per-expert MLP weights. These are exactly the real nn.Linear layers we want to quantize. So without the filter, essentially every encoder and decoder attention/MLP linear is excluded from quantization, it's a bug. The filter restricts the auto-derived skip list back to "embed/lm_head/norm", which is what get_keys_to_not_convert was meant to return.
Do you have better idea (need to be simple) to fix this issue?
|
could you leverage auto-round-best tune a model and upload to intel space, thanks! |
yes, WIP. |
Signed-off-by: lvliang-intel <liang1.lv@intel.com>
|
/azp run Unit-Test-CUDA-AutoRound |
|
Azure Pipelines successfully started running 1 pipeline(s). |
…upport_diffusiongemma
…/auto-round into lvl/support_diffusiongemma
|
/azp run Unit-Test-CUDA-AutoRound |
|
Azure Pipelines successfully started running 1 pipeline(s). |
|
/azp run Unit-Test-CUDA-AutoRound |
|
Azure Pipelines successfully started running 1 pipeline(s). |
…upport_diffusiongemma Signed-off-by: lvliang-intel <liang1.lv@intel.com>
|
/azp run Unit-Test-CUDA-AutoRound |
|
Azure Pipelines successfully started running 1 pipeline(s). |
| name for name in modules_to_not_convert if any(key in name for key in _DEFAULT_SKIP_KEYWORDS) | ||
| ] | ||
|
|
||
| if modules_to_not_convert: |
There was a problem hiding this comment.
module_to_not_convert is mainly used for the AWQ format. In vLLM and SGLang, we only parse extra_config.
Besides, if this is caused by the tied-weight key issue, how about handling tied-weight keys explicitly instead?
|
|
||
| if modules_to_not_convert: | ||
| # Pre-compile patterns once instead of recompiling them for every layer name. | ||
| compiled_patterns = [re.compile(n) for n in modules_to_not_convert] |
There was a problem hiding this comment.
@WeiweiZhang1 @n1ck-guo please review this change. I have some concerns that it might be risky.
There was a problem hiding this comment.
From the code, the risk is not significant. The only issue is that the impact is large. Could we limit it to a controllable range using methods like patching, or perhaps control it through parameters?
There was a problem hiding this comment.
If the user mark some layers as non-quantized, there should be no entries in module_not_convert when the format is AutoRound instead of AWQ. In AutoRound, we only use extra_config to indicate whether the lm-head is quantized or to mark certain layers as not quantized.
So I think handling this logic via module_not_convert might be problematic.
There was a problem hiding this comment.
Agree with this point. I will change the code as below.
# AWQ format: exclude specified modules.
if "awq" in (getattr(quantization_config, "quant_method", "") or "").lower():
extra_config = skip_not_convert_modules(model, quantization_config, layer_names, extra_config)
Signed-off-by: lvliang-intel <liang1.lv@intel.com>
Signed-off-by: lvliang-intel <liang1.lv@intel.com>
…/auto-round into lvl/support_diffusiongemma
4 participants
Description
Add quantization support for DiffusionGemma (transformers ≥ 5.11.0). The model declares encoder/decoder MoE gate_up_proj / down_proj as tied across stacks; AutoRound's MoE-unfuse step splits these fused 3D parameters into per-expert linears, which leaves the declared _tied_weights_keys referring to nonexistent parameter names and breaks downstream loading/export.
Type of Change
New feature
Related Issues
None
Checklist Before Submitting
/azp run Unit-Test-CUDA-AutoRound.